package adkr

import (
    "ebbflow/pkg/utils"
    "bytes"
    "encoding/json"
    "log"
)

// ReliableBroadcastDual implements a dual/dual_cp variant of RBC used in HBACSS
// mode: "dual" or "dual_cp" (checkpoint assist)
// On deliver, it writes the reconstructed payload bytes to outCh.
type dualValPayload struct {
    RootHash []byte   `json:"root_hash"`
    Branch   [][]byte `json:"branch"`
    Stripe   []byte   `json:"stripe"`
    TotalLen int      `json:"total_len"`
}

func ReliableBroadcastDualMode(
    sid string,
    pid, n, f int,
    leader int,
    mode string,
    inCh <-chan []byte,
    outCh chan<- []byte,
    msgInCh <-chan Message,
    msgOutCh chan<- Message,
) {
    // thresholds
    k := n - 2*f
    echoThreshold := n - f
    readyThreshold := f + 1
    outputThreshold := 2*f + 1

    send := func(to int, m RBCMessage) {
        payload, _ := json.Marshal(m)
        msgOutCh <- Message{Tag: ADKR_MVBA, Round: leader, To: to, Payload: payload}
    }
    bcast := func(m RBCMessage) {
        payload, _ := json.Marshal(m)
        msgOutCh <- Message{Tag: ADKR_MVBA, Round: leader, To: -1, Payload: payload}
    }

    // Leader encodes and sends VAL stripes with total length
    var totalLen int
    if pid == leader {
        data := <-inCh
        totalLen = len(data)
        stripes, err := utils.Encode(k, n, data)
        if err != nil { log.Printf("RBC-DUAL %s encode error: %v", sid, err); return }
        mt, root := utils.MerkleTree(stripes)
        for i := 0; i < n; i++ {
            branch, _ := utils.GetMerkleBranch(i, mt)
            val := dualValPayload{RootHash: root, Branch: branch, Stripe: stripes[i], TotalLen: totalLen}
            b, _ := json.Marshal(val)
            send(i, RBCMessage{Tag: RBCValTag, Payload: b})
        }
        log.Printf("RBC-DUAL %s leader sent VAL root=%x", sid, root[:4])
    }

    var rootRef []byte
    stripes := make(map[string][][]byte)
    lens := make(map[string]int)
    echoCnt := make(map[string]int)
    readySet := make(map[string]map[int]struct{})
    readySent := false
    checkpointed := false

    for msg := range msgInCh {
        if msg.Tag != ADKR_MVBA || msg.Round != leader { continue }
        var m RBCMessage
        if err := json.Unmarshal(msg.Payload, &m); err != nil { continue }

        switch m.Tag {
        case RBCValTag:
            if msg.Sender != leader { continue }
            var val dualValPayload
            if json.Unmarshal(m.Payload, &val) != nil { continue }
            if !utils.MerkleVerify(n, val.Stripe, val.RootHash, val.Branch, pid) { continue }
            if rootRef == nil { rootRef = val.RootHash }
            lens[string(val.RootHash)] = val.TotalLen
            // Immediately ECHO the stripe we received
            bcast(RBCMessage{Tag: RBCEchoTag, Payload: m.Payload})

        case RBCEchoTag:
            var echo dualValPayload
            if json.Unmarshal(m.Payload, &echo) != nil { continue }
            if !utils.MerkleVerify(n, echo.Stripe, echo.RootHash, echo.Branch, msg.Sender) { continue }
            key := string(echo.RootHash)
            if _, ok := stripes[key]; !ok { stripes[key] = make([][]byte, n); readySet[key] = make(map[int]struct{}) }
            if stripes[key][msg.Sender] != nil { continue }
            stripes[key][msg.Sender] = echo.Stripe
            if _, ok := lens[key]; !ok { lens[key] = echo.TotalLen }
            echoCnt[key]++
            if echoCnt[key] >= echoThreshold && !readySent {
                readySent = true
                rp, _ := json.Marshal(RBCReadyPayload{RootHash: echo.RootHash})
                bcast(RBCMessage{Tag: RBCReadyTag, Payload: rp})
            }
            if mode == "dual_cp" && !checkpointed {
                have := 0
                for _, s := range stripes[key] { if s != nil { have++ } }
                if have >= k { checkpointed = true }
            }

        case RBCReadyTag:
            var rdy RBCReadyPayload
            if json.Unmarshal(m.Payload, &rdy) != nil { continue }
            key := string(rdy.RootHash)
            if _, ok := readySet[key]; !ok { readySet[key] = make(map[int]struct{}) }
            if _, ok := readySet[key][msg.Sender]; ok { continue }
            readySet[key][msg.Sender] = struct{}{}
            if len(readySet[key]) >= readyThreshold && !readySent {
                readySent = true
                bcast(RBCMessage{Tag: RBCReadyTag, Payload: m.Payload})
            }
            if len(readySet[key]) >= outputThreshold {
                have := 0
                for _, s := range stripes[key] { if s != nil { have++ } }
                if have >= k {
                    decoded, err := utils.Decode(k, n, stripes[key])
                    if err != nil { log.Printf("RBC-DUAL %s decode error: %v", sid, err); return }
                    // Trim to total length if provided
                    if L, ok := lens[key]; ok && L >= 0 && L <= len(decoded) { decoded = decoded[:L] }
                    re, _ := utils.Encode(k, n, decoded)
                    _, reRoot := utils.MerkleTree(re)
                    if !bytes.Equal(reRoot, rdy.RootHash) { log.Printf("RBC-DUAL %s root mismatch", sid); return }
                    outCh <- decoded
                    return
                }
            }
        }
    }
}

// Thin wrappers
func ReliableBroadcastDual(sid string, pid, n, f, dealer int, inCh <-chan []byte, outCh chan<- []byte, msgIn <-chan Message, msgOut chan<- Message) {
    ReliableBroadcastDualMode(sid, pid, n, f, dealer, "dual", inCh, outCh, msgIn, msgOut)
}

func ReliableBroadcastDualCP(sid string, pid, n, f, dealer int, inCh <-chan []byte, outCh chan<- []byte, msgIn <-chan Message, msgOut chan<- Message) {
    ReliableBroadcastDualMode(sid, pid, n, f, dealer, "dual_cp", inCh, outCh, msgIn, msgOut)
}

// StartRBCWithMode picks RBC variant based on GetRBCMode()
func StartRBCWithMode(sid string, pid, n, f, dealer int, inCh <-chan []byte, outCh chan<- []byte, msgIn <-chan Message, msgOut chan<- Message) {
    switch GetRBCMode() {
    case "dual":
        ReliableBroadcastDual(sid, pid, n, f, dealer, inCh, outCh, msgIn, msgOut)
    case "dual_cp":
        ReliableBroadcastDualCP(sid, pid, n, f, dealer, inCh, outCh, msgIn, msgOut)
    default:
        ReliableBroadcast(sid, pid, n, f, dealer, inCh, outCh, msgIn, msgOut)
    }
}
